from .encoder import *


class DecoderLayer(nn.Module):
    def __init__(self, d_model, mask_self_attn, src_attn, feed_forward, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.mask_self_attn = mask_self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward
        self.dropout = dropout
        self.layers = clones(SubLayerConnection(self.d_model, self.dropout), 3)

    def forward(self, data, encoder_output, source_mask, target_mask):
        result = self.layers[0](
            data,
            lambda x: self.mask_self_attn(query=x, key=x, value=x, mask=target_mask),
        )
        result = self.layers[1](
            result,
            lambda x: self.src_attn(
                query=x, key=encoder_output, value=encoder_output, mask=source_mask
            ),
        )
        result = self.layers[2](result, lambda x: self.feed_forward(x))
        return result


class Decoder(nn.Module):
    def __init__(self, layer: DecoderLayer, N):
        super().__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.d_model)

    def forward(self, data, encoder_output, source_mask, target_mask):
        for layer in self.layers:
            data = layer(data, encoder_output, source_mask, target_mask)
        return self.norm(data)
